package com.github.manx98.okdownloader.utils;

import com.github.manx98.okdownloader.common.DownloadErrorLevel;
import com.github.manx98.okdownloader.exception.ReadDataFromResponseErrorException;
import com.github.manx98.okdownloader.exception.RequestFailException;
import com.github.manx98.okdownloader.exception.WriteDataErrorException;
import com.github.rholder.retry.RetryException;
import okhttp3.Call;
import okhttp3.Request;
import okhttp3.Response;
import okio.BufferedSource;
import com.github.manx98.okdownloader.common.TaskStatusType;
import com.github.manx98.okdownloader.event.DownloadTaskEventHandler;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;

/**
 * <b>下载执行器</b>
 * <br/>
 * 此类实现了<b>Runnable接口</b>，但是同一个实例是<b>不允许被多个线程同时执行</b>。
 * 但是如果在下载过程中出现问题题，导致整个下载流程终止时，这个实例仍可以被再次运行，恢复下载。
 *
 * @author : manx98
 * @date : 2021/5/27
 */
public class OkDownloader implements Runnable, AutoCloseable {
    private final DownloadTask downloadTask;
    private final ThreadPoolExecutor pool;
    private final AtomicBoolean alive = new AtomicBoolean(true);
    private final AtomicBoolean paused = new AtomicBoolean(false);
    private final DownloadTaskStatus downloadTaskStatus;
    private final DataSaveHandler dataSaveHandler;
    private final DownloadTaskEventHandler downloadEventHandler;
    private final ProgressHandler progressHandler;
    private final long minSplitX2;
    private final List<DownloadBlockInfo> downloadBlockInfos;
    private final Logger logger = LoggerFactory.getLogger(OkDownloader.class);
    private final LinkedBlockingDeque<DownloadBlockInfo> tasksQueue = new LinkedBlockingDeque<>();
    private final ThreadLocal<ByteBuffer> byteBufferThreadLocal;
    private CountDownLatch workerCountDown;

    /**
     * 创建下载任务执行器
     *
     * @param downloadTask       下载任务对象
     * @param threadPoolExecutor 下载执行任务时使用的线程池,请自行确保可用线程大于下载任务的线程
     * @throws IOException 读取进度文件失败时抛出
     */
    public OkDownloader(DownloadTask downloadTask, ThreadPoolExecutor threadPoolExecutor) throws IOException {
        this.pool = threadPoolExecutor;
        this.dataSaveHandler = downloadTask.getDataSaveHandler();
        this.downloadEventHandler = downloadTask.getDownloadEventHandler();
        this.progressHandler = downloadTask.getProgressHandler();
        this.downloadTask = downloadTask;
        this.minSplitX2 = Math.min(progressHandler.getMinSplitSize(), progressHandler.getMaxSplitSize()) << 1;
        this.downloadBlockInfos = this.progressHandler.getDownloadBlockProgressInfos();
        this.downloadTaskStatus = downloadTask.getDownloadTaskStatus();
        this.dataSaveHandler.fileShrink(this.downloadTaskStatus.getFileSize());
        byteBufferThreadLocal = ThreadLocal.withInitial(() -> ByteBuffer.allocateDirect(downloadTask.getBufferSize()));
    }

    @Override
    public void run() {
        alive.set(true);
        paused.set(false);
        tasksQueue.clear();
        downloadEventHandler.onStart(downloadTask);
        try {
            this.downloadTaskStatus.setStatus(TaskStatusType.RUNNING);
            init();
            intervalUpdateSpeed();
            if (paused.get()) {
                this.downloadTaskStatus.setStatus(TaskStatusType.PAUSE);
                this.downloadEventHandler.onPause(this.downloadTask);
            } else {
                if (alive.get() && this.downloadTaskStatus.getCompleteSize() >= this.downloadTaskStatus.getFileSize()) {
                    this.downloadTaskStatus.setStatus(TaskStatusType.FINISHED);
                    downloadEventHandler.onComplete(downloadTask);
                    return;
                }
                downloadEventHandler.onFailed(downloadTask);
                this.downloadTaskStatus.setStatus(TaskStatusType.ERROR);
            }
        } finally {
            alive.set(false);
            downloadTaskStatus.setSpeed();
            downloadEventHandler.onFinally(this);
        }
    }

    /**
     * 周期更新下载速度
     */
    private void intervalUpdateSpeed() {
        try {
            while (workerCountDown.getCount() > 0) {
                this.updateCompleteSize();
                TimeUnit.SECONDS.sleep(1);
            }
        } catch (InterruptedException e) {
            logger.info("自动更新下载实时速度被中断:{}", e.getMessage());
            Thread.currentThread().interrupt();
        } finally {
            this.updateCompleteSize();
        }
    }

    /**
     * 更新任务完成长度
     */
    private void updateCompleteSize() {
        long[] uncompletedSize = new long[]{0};
        this.downloadBlockInfos.forEach(downloadBlockInfo -> {
            long start = downloadBlockInfo.getStart();
            long end = downloadBlockInfo.getEnd();
            if (start <= end) {
                uncompletedSize[0] += end - start + 1;
            }
            if (!this.alive.get()) {
                final Call call = downloadBlockInfo.getCall();
                if (call != null) {
                    call.cancel();
                }
            }
        });
        downloadTaskStatus.updateUnCompleteSize(uncompletedSize[0]);
    }

    /**
     * 初始化下载任务
     */
    private void init() {
        if (downloadBlockInfos.isEmpty()) {
            this.workerCountDown = new CountDownLatch(0);
            return;
        }
        tasksQueue.addAll(downloadBlockInfos);
        final List<DownloadLink> downloadLinks = downloadTask.getDownloadLinks();
        int[] w = new int[downloadLinks.size()];
        int i = 0;
        int maxWorkers = this.progressHandler.getMaxWorkers();
        for (DownloadLink downloadLink : downloadLinks) {
            w[i++] = downloadLink.getMaxWorkers();
        }
        boolean hasMore = true;
        List<DownloadLink> workers = new LinkedList<>();
        while (hasMore) {
            int f = 0;
            i = 0;
            for (DownloadLink downloadLink : downloadLinks) {
                if (w[i] > 0) {
                    maxWorkers -= 1;
                    workers.add(downloadLink);
                    w[i] -= 1;
                    f += 1;
                }
                if (maxWorkers <= 0) {
                    break;
                }
                ++i;
            }
            hasMore = f > 0 && maxWorkers > 0;
        }
        this.workerCountDown = new CountDownLatch(workers.size());
        workers.forEach(this::addExecutor);
    }

    /**
     * 暂停正在执行的下载任务(非阻塞阻塞)
     */
    public synchronized void pause() {
        if (alive.get()) {
            paused.set(true);
            alive.set(false);
        }
    }

    /**
     * 暂停正在执行的下载任务(塞阻塞)
     */
    public synchronized void pauseSync() throws InterruptedException {
        pause();
        workerCountDown.await();
    }

    /**
     * 执行下载
     *
     * @param downloadBlockInfoWithLink 区间信息
     */
    private void executeTask(DownloadBlockInfoWithLink downloadBlockInfoWithLink) throws RequestFailException, WriteDataErrorException, ReadDataFromResponseErrorException {
        if (!this.alive.get()) {
            return;
        }
        final DownloadLink link = downloadBlockInfoWithLink.getLink();
        Request downloadRequest = link.getDownloadRequest(downloadBlockInfoWithLink.getStart(), downloadBlockInfoWithLink.getEnd());
        final Call call = downloadTask.getClient().newCall(downloadRequest);
        downloadBlockInfoWithLink.getDownloadBlockProgressInfo().setCall(call);
        try (Response execute = call.execute()) {
            if (execute.isSuccessful()) {
                logger.debug("链接建立成功\n├─[数据段]:{} - {}\n├─[下载链接]:{}\n└─[HTTP 状态码]:{}"
                        , downloadBlockInfoWithLink.getStart()
                        , downloadBlockInfoWithLink.getEnd()
                        , link.getUrl()
                        , execute.code()
                );
                try (BufferedSource inputStream = Objects.requireNonNull(execute.body()).source()) {
                    saveData(inputStream, downloadBlockInfoWithLink);
                }
            } else {
                logger.warn("服务器响应异常\n├─[数据段]:{} - {}\n├─[下载链接]:{}\n└─[HTTP 状态码]:{}"
                        , downloadBlockInfoWithLink.getStart()
                        , downloadBlockInfoWithLink.getEnd()
                        , link.getUrl()
                        , execute.code()
                );
                throw new RequestFailException(execute.code());
            }
        } catch (IOException e) {
            if (!call.isCanceled()) {
                logger.warn("建立下载链接失败\n├─[数据段]:{} - {}\n├─[下载链接]:{}\n└─[异常]:{}"
                        , downloadBlockInfoWithLink.getStart()
                        , downloadBlockInfoWithLink.getEnd()
                        , link.getUrl()
                        , e.getMessage()
                );
                throw new RequestFailException("建立下载链接失败:" + e.getMessage(), e);
            }
        } finally {
            logger.debug("链接建立关闭\n├─[下载链接]:{}\n└─[数据段]:{} - {}"
                    , link.getUrl()
                    , downloadBlockInfoWithLink.getStart()
                    , downloadBlockInfoWithLink.getEnd()
            );
            downloadBlockInfoWithLink.getDownloadBlockProgressInfo().setCall(null);
        }
    }

    /**
     * 写入数据
     *
     * @param inputStream               数据流
     * @param downloadBlockInfoWithLink 下载块
     */
    private void saveData(BufferedSource inputStream, DownloadBlockInfoWithLink downloadBlockInfoWithLink) throws WriteDataErrorException, ReadDataFromResponseErrorException {
        DownloadBlockInfo downloadBlockInfo = downloadBlockInfoWithLink.getDownloadBlockProgressInfo();
        final ByteBuffer data = byteBufferThreadLocal.get();
        int n;
        try {
            while (alive.get() && (n = inputStream.read(data)) > 0) {
                try {
                    dataSaveHandler.save(data, downloadBlockInfo.getStart());
                    downloadBlockInfo.complete(n);
                    progressHandler.updateStartProgress(downloadBlockInfo);
                    if (downloadBlockInfo.getStart() > downloadBlockInfo.getEnd()) {
                        downloadBlockInfoWithLink.setCode(DownloadErrorLevel.OK);
                        return;
                    }
                } catch (IOException e) {
                    downloadEventHandler.onError(downloadBlockInfoWithLink, e);
                    logger.error("下载数据写入失败\n├─[下载链接]:{}\n├─[数据段]:{} - {}\n└─[异常]:{}", downloadBlockInfoWithLink.getLink().getUrl(), downloadBlockInfo.getStart(), downloadBlockInfo.getEnd(), e.getMessage());
                    throw new WriteDataErrorException("下载数据写入失败:" + e.getMessage(), e);
                }
            }
        } catch (IOException e) {
            throw new ReadDataFromResponseErrorException(e);
        }
    }

    /**
     * 查找剩余大于2倍最小分割的下载的最长的下载任务
     *
     * @return 返回下载块，当无法找到合适的块时返回null
     */
    private DownloadBlockInfo getLongestTask() {
        DownloadBlockInfo[] best = new DownloadBlockInfo[]{null};
        long[] bestUncompletedSize = new long[]{0L};
        if (!downloadBlockInfos.isEmpty()) {
            downloadBlockInfos.forEach(progressInfo -> {
                if (progressInfo.getCode() != DownloadErrorLevel.OK) {
                    long uncompletedSize = progressInfo.getEnd() - progressInfo.getStart() + 1;
                    if (uncompletedSize >= minSplitX2 && uncompletedSize > bestUncompletedSize[0]) {
                        best[0] = progressInfo;
                    }
                }
            });
        }
        return best[0];
    }

    /**
     * 动态分割现有下载任务
     *
     * @param downloadBlockInfo 需要分割任务
     * @throws IOException 进度写入失败
     */
    private synchronized void trySplitTask(DownloadBlockInfo downloadBlockInfo) throws IOException {
        while (this.alive.get() && this.downloadTaskStatus.getCompleteSize() < this.downloadTaskStatus.getFileSize()) {
            logger.debug("开始重划分现有未完成数据块！");
            final DownloadBlockInfo best = getLongestTask();
            if (best == null) {
                logger.debug("无法划分现有未完成数据块！");
                break;
            } else {
                long start = best.getStart();
                long end = best.getEnd();
                long mid = start + (end - start) / 2;
                if (mid + 1 <= end) {
                    downloadBlockInfo.setStart(mid + 1);
                    downloadBlockInfo.setEnd(end);
                    best.setEnd(mid);
                    progressHandler.updateProgress(downloadBlockInfo);
                    progressHandler.updateEndProgress(best);
                    addTask(downloadBlockInfo);
                    logger.debug("开始重划分现有未完成数据块成功\n└─[拆分范围]:{}-{}", downloadBlockInfo.getStart(), downloadBlockInfo.getEnd());
                    return;
                }
            }
        }
        this.downloadBlockInfos.remove(downloadBlockInfo);
    }

    /**
     * 添加下载块
     *
     * @param progressInfo 下载块
     */
    private void addTask(DownloadBlockInfo progressInfo) {
        if (alive.get()) {
            progressInfo.setCode(DownloadErrorLevel.WARNING);
            tasksQueue.add(progressInfo);
        }
    }

    /**
     * 添加执行核心
     * 以链接为单位创建执行器
     *
     * @param downloadLink 下载链接
     */
    private void addExecutor(DownloadLink downloadLink) {
        pool.execute(() -> {
            try {
                while (alive.get()) {
                    DownloadBlockInfo downloadBlockInfo = tasksQueue.poll(500, TimeUnit.MILLISECONDS);
                    if (downloadBlockInfo == null) {
                        if (this.downloadTaskStatus.getCompleteSize() >= this.downloadTaskStatus.getFileSize() || this.downloadBlockInfos.isEmpty()) {
                            return;
                        }
                    } else {
                        if (alive.get() && !executor(downloadLink, downloadBlockInfo)) {
                            return;
                        }
                    }
                }
            } catch (InterruptedException e) {
                logger.error("下载核心线程被中断:{}", e.getMessage());
                Thread.currentThread().interrupt();
            } finally {
                this.workerCountDown.countDown();
                progressHandler.onThreadTermination();
                byteBufferThreadLocal.remove();
            }
        });
    }

    /**
     * 下载任务执行单元
     *
     * @param downloadLink      下载链接
     * @param downloadBlockInfo 任务块
     * @return 是否执行成功
     */
    private boolean executor(DownloadLink downloadLink, DownloadBlockInfo downloadBlockInfo) {
        downloadBlockInfo.setCode(DownloadErrorLevel.RUNNING);
        downloadTaskStatus.threadCount.increment();
        DownloadBlockInfoWithLink downloadBlockInfoWithLink = new DownloadBlockInfoWithLink(downloadBlockInfo, downloadLink);
        try {
            executeWithRetryer(downloadBlockInfoWithLink);
            downloadBlockInfo.setCode(DownloadErrorLevel.OK);
            if (tasksQueue.isEmpty() && this.alive.get() && this.downloadTaskStatus.getCompleteSize() < this.downloadTaskStatus.getFileSize()) {
                long startTime = System.currentTimeMillis();
                logger.debug("等待划分任务：{}", downloadBlockInfo);
                trySplitTask(downloadBlockInfo);
                logger.debug("划分任务结束,耗时{}ms：{}", System.currentTimeMillis() - startTime, downloadBlockInfo);
            }
            return true;
        } catch (IOException e) {
            this.downloadTaskStatus.setException(e);
            downloadEventHandler.onError(downloadBlockInfoWithLink, e);
            alive.set(false);
            logger.error("进度划分写入出现致命错误,下载任务即将被终止!\n├─[最大线程]:{}├─[下载链接]:{}\n├─[数据切片范围]:{} - {}\n└─[异常]:{}", downloadLink.getMaxWorkers(), downloadLink.getUrl(), downloadBlockInfo.getStart(), downloadBlockInfo.getEnd(), e.getMessage());
            return false;
        } catch (RetryException e) {
            logger.warn("下载块重试失败,将被重新加入任务队列\n├─[最大线程]:{}├─[下载链接]:{}\n├─[数据切片范围]:{} - {}\n└─[异常]:{}", downloadLink.getMaxWorkers(), downloadLink.getUrl(), downloadBlockInfo.getStart(), downloadBlockInfo.getEnd(), e.getMessage());
            this.downloadTaskStatus.setException(e);
            downloadEventHandler.onError(downloadBlockInfoWithLink, e);
            this.downloadTaskStatus.setException(e.getLastFailedAttempt().getExceptionCause());
            addTask(downloadBlockInfo);
            return false;
        } catch (ExecutionException e) {
            this.downloadTaskStatus.setException(e);
            downloadEventHandler.onError(downloadBlockInfoWithLink, e);
            alive.set(false);
            logger.error("下载块出现不可恢复异常,下载任务即将被终止！\n├─[最大线程]:{}\n├─[下载链接]:{}\n├─[数据切片范围]:{} - {}\n└─[异常]:{}", downloadLink.getMaxWorkers(), downloadLink.getUrl(), downloadBlockInfo.getStart(), downloadBlockInfo.getEnd(), e.getMessage());
            return false;
        } finally {
            downloadTaskStatus.threadCount.decrement();
        }
    }

    /**
     * 使用重试机制运行下载任务
     *
     * @param downloadBlockInfoWithLink 下载任务块
     */
    private void executeWithRetryer(DownloadBlockInfoWithLink downloadBlockInfoWithLink) throws ExecutionException, RetryException {
        if (alive.get()) {
            downloadBlockInfoWithLink.getLink().getRetryer().call(() -> {
                logger.debug("开始执行下载块：{}", downloadBlockInfoWithLink);
                executeTask(downloadBlockInfoWithLink);
                logger.debug("下载块执行成功：{}", downloadBlockInfoWithLink);
                return null;
            });
        }
    }

    /**
     * 获取下载任务状态
     *
     * @return 任务状态
     */
    public DownloadTaskStatus getDownloadTaskStatus() {
        return downloadTaskStatus;
    }

    @Override
    public void close() throws Exception {
        try {
            dataSaveHandler.close();
        } finally {
            progressHandler.close();
        }
    }
}
